# coding=utf-8
# Chern-C
# Copyright (C) 2024 NJU-SAML
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from asyncio import AbstractEventLoop
import json
import logging
import logging.handlers
import os
import platform
import re
import sys
from typing import AsyncGenerator, Dict, Generator, Optional, Union
import warnings
from chernc.constants import DEFAULT_LOGGER_DIR, DEFAULT_LOGGER_FILE_NAME

FILE_HANDLERS: Dict[str, logging.FileHandler] = {}


# Define a custom formatter without color codes
class NoColorFormatter(logging.Formatter):
    color_pattern = re.compile(r"\x1b[^m]*m")  # Regex pattern to match color codes

    def format(self, record):
        message = super().format(record)
        # Remove color codes from the log message
        message = self.color_pattern.sub("", message)
        return message


FORMATTER = logging.Formatter(
    fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

NOCOLOR_FORMATTER = NoColorFormatter(
    fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

SIMPLE_FORMATTER = NoColorFormatter(
    fmt="%(asctime)s | %(levelname)s | %(name)s |\n%(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)


def build_simple_file_logger(logger_name, logger_filename: str) -> logging.Logger:
    # Get logger
    logger = logging.getLogger(logger_name)
    logger.setLevel(logging.DEBUG)

    # Add a file handler
    dir_path = os.path.dirname(logger_filename)
    os.makedirs(dir_path, exist_ok=True)
    
    handler = logging.FileHandler(logger_filename, mode="w", encoding="utf-8")
    handler.setFormatter(SIMPLE_FORMATTER)
    logger.addHandler(handler)

    return logger


def build_logger(
    logger_name,
    logger_filename: Optional[str] = DEFAULT_LOGGER_FILE_NAME,
    logger_dir: Optional[str] = DEFAULT_LOGGER_DIR,
    log_all: bool = True,
) -> logging.Logger:
    if logger_dir is None:
        logger_dir = "./"
    if logger_filename is None:
        logger_filename = DEFAULT_LOGGER_FILE_NAME

    global FILE_HANDLERS, NOCOLOR_FORMATTER, FORMATTER

    # Get logger
    logger = logging.getLogger(logger_name)
    logger.setLevel(logging.DEBUG)

    # Add a file handler for all loggers
    if logger_filename is not None:
        if logger_filename not in FILE_HANDLERS:
            os.makedirs(logger_dir, exist_ok=True)
            filename = os.path.join(logger_dir, logger_filename)
            handler = logging.handlers.TimedRotatingFileHandler(filename, when="H", utc=True, encoding="utf-8")
            handler.setFormatter(NOCOLOR_FORMATTER)
            handler.namer = lambda name: name.replace(".log", "") + ".log"
            FILE_HANDLERS[logger_filename] = handler
        else:
            handler = FILE_HANDLERS[logger_filename]

        if log_all:
            for name, item in logging.root.manager.loggerDict.items():
                if isinstance(item, logging.Logger):
                    if handler not in item.handlers:
                        item.addHandler(handler)
        else:
            logger.addHandler(handler)

    return logger


def release_logger(logger: Union[logging.Logger, str]) -> None:
    if isinstance(logger, str):
        logger = logging.getLogger(logger)
    for handler in logger.handlers:
        if isinstance(handler, logging.handlers.TimedRotatingFileHandler):
            continue
        if isinstance(handler, logging.FileHandler):
            handler.close()
    logger.handlers.clear()
    
    del logging.root.manager.loggerDict[logger.name]
    
class StreamToLogger(object):
    """
    Fake file-like stream object that redirects writes to a logger instance.
    """

    def __init__(self, logger, log_level=logging.INFO):
        self.terminal = sys.stdout
        self.logger = logger
        self.log_level = log_level
        self.linebuf = ""

    def __getattr__(self, attr):
        try:
            attr_value = getattr(self.terminal, attr)
        except:
            return None
        return attr_value

    def write(self, buf):
        temp_linebuf = self.linebuf + buf
        self.linebuf = ""
        for line in temp_linebuf.splitlines(True):
            # From the io.TextIOWrapper docs:
            #   On output, if newline is None, any '\n' characters written
            #   are translated to the system default line separator.
            # By default sys.stdout.write() expects '\n' newlines and then
            # translates them so this is still cross platform.
            if line[-1] == "\n":
                encoded_message = line.encode("utf-8", "ignore").decode("utf-8")
                self.logger.log(self.log_level, encoded_message.rstrip())
            else:
                self.linebuf += line

    def flush(self):
        if self.linebuf != "":
            encoded_message = self.linebuf.encode("utf-8", "ignore").decode("utf-8")
            self.logger.log(self.log_level, encoded_message.rstrip())
        self.linebuf = ""


# Set the format of root handlers
root_logger = logging.getLogger()

console_handler = logging.StreamHandler()
console_handler.setFormatter(FORMATTER)

# the message printed to console are limited
console_handler.setLevel(logging.INFO)
root_logger.addHandler(console_handler)
# we allow any levels of messages
root_logger.setLevel(logging.DEBUG)


# Set root logger level
# if sys.version_info[1] >= 9:
#     # This is for windows
#     logging.basicConfig(level=logging.INFO, encoding="utf-8")
# else:
#     if platform.system() == "Windows":
#         warnings.warn("If you are running on Windows, " "we recommend you use Python >= 3.9 for UTF-8 encoding.")
#     logging.basicConfig(level=logging.INFO)

# Set all logger formatter
# for handler in logging.getLogger().handlers:
#     if isinstance(handler, logging.FileHandler):
#         continue
#     if handler.formatter != formatter:
#         handler.setFormatter(formatter)

# Redirect stdout and stderr to loggers
if not isinstance(sys.stdout, StreamToLogger):
    stdout_logger = logging.getLogger("stdout")
    stdout_logger.setLevel(logging.DEBUG)
    sl = StreamToLogger(stdout_logger, logging.INFO)
    sys.stdout = sl

if not isinstance(sys.stderr, StreamToLogger):
    stderr_logger = logging.getLogger("stderr")
    stderr_logger.setLevel(logging.ERROR)
    sl = StreamToLogger(stderr_logger, logging.ERROR)
    sys.stderr = sl


def iter_over_async(async_gen: AsyncGenerator, event_loop: AbstractEventLoop) -> Generator:
    """
    Convert async generator to sync generator

    :param async_gen: the AsyncGenerator to convert
    :param event_loop: the event loop to run on
    :returns: Sync generator
    """
    ait = async_gen.__aiter__()

    async def get_next():
        try:
            obj = await ait.__anext__()
            return False, obj
        except StopAsyncIteration:
            return True, None

    while True:
        done, obj = event_loop.run_until_complete(get_next())
        if done:
            break
        yield obj
